import os
import sys
import subprocess
from typing import Dict, List, Union
import re

import pandas as pd
import json
from pandas import Timestamp

from weclone.data.clean.strategies import LLMCleaningStrategy
from weclone.utils.config import load_config
from weclone.utils.log import logger
from weclone.data.models import ChatMessage, CutMessage, skip_type_list, QaPair
from weclone.data.strategies import TimeWindowStrategy, LLMStrategy


class DataProcessor:
    def __init__(self):
        self.config = load_config(arg_type="make_dataset")
        self.csv_folder = "./dataset/csv"
        self.system_prompt = self.config["default_system"]
        self.cut_type_list = [
            "图片",
            "视频",
            "合并转发的聊天记录",
            "语音",
            "(分享)音乐",
            "(分享)卡片式链接",
            "(分享)笔记",
            "(分享)小程序",
            "(分享)收藏夹",
            "(分享)小说(猜)",
            "(分享)视频号名片",
            "(分享)视频号视频",
            "粘贴的文本",  # 无法解析的分享链接
        ]

        # blocked_words
        config_blocked_words = self.config.get("blocked_words", [])
        file_blocked_words = []
        try:
            with open("./dataset/blocked_words.json", encoding="utf-8") as f:
                file_blocked_words = json.load(f).get("blocked_words", [])
        except (FileNotFoundError, json.JSONDecodeError):
            pass

        self.blocked_words = list(set(config_blocked_words + file_blocked_words))
        logger.info(f"聊天记录禁用词: {self.blocked_words}")

        if self.config.get("clean_dataset", {}).get("enable_clean", False) and self.config.get(
            "prompt_with_history", False
        ):
            logger.warning("开启 prompt_with_history 不支持 clean_dataset 功能")
            exit()

        if self.config["single_combine_strategy"] == "time_window":
            self.single_combine_strategy = TimeWindowStrategy(
                time_window=self.config["single_combine_time_window"] * 60,
                is_single_chat=True,
            )
        elif self.config["single_combine_strategy"] == "llm":
            self.single_combine_strategy = LLMStrategy(
                is_single_chat=True,
            )

        if self.config["qa_match_strategy"] == "time_window":
            self.qa_match_strategy = TimeWindowStrategy(
                time_window=self.config["qa_match_time_window"] * 60,
                is_single_chat=False,
            )
        elif self.config["qa_match_strategy"] == "llm":
            self.qa_match_strategy = LLMStrategy(is_single_chat=False)

        if self.config.get("clean_dataset", {}).get("enable_clean", False):
            if self.config.get("clean_dataset", {}).get("clean_strategy", "llm") == "llm":
                self.clean_strategy = LLMCleaningStrategy(make_dataset_config=self.config)
        self.c = self.config

    def main(self):
        if not os.path.exists(self.csv_folder) or not os.listdir(self.csv_folder):
            logger.error(f"错误：目录 '{self.csv_folder}' 不存在或为空，请检查路径并确保其中包含 CSV 聊天数据文件。")
            return

        csv_files = self.get_csv_files()
        message_list: List[ChatMessage] = []
        for csv_file in csv_files:
            chat_messages = self.load_csv(csv_file)
            message_list.extend(self.group_consecutive_messages(messages=chat_messages))
            # self.process_by_msgtype(chat_message)
        qa_res = self.match_qa(message_list)
        if self.c["prompt_with_history"]:
            qa_res = self.add_history_to_qa(qa_res)
        else:
            qa_res = [item for item in qa_res if isinstance(item, QaPair)]

        if self.c.get("clean_dataset", {}).get("enable_clean", False):
            self.clean_strategy.judge(qa_res)
            qa_res = self.clean_strategy.clean(qa_res)
        self.save_result(qa_res)
        self._execute_length_cdf_script()

        logger.success(f"聊天记录处理成功，共{len(qa_res)}条，保存到 ./dataset/res_csv/sft/sft-my.json")

    def _execute_length_cdf_script(self):
        """执行 length_cdf.py 脚本来计算cutoff_len。"""
        try:
            python_executable = sys.executable
            # 脚本路径是相对于项目根目录的
            script_path = os.path.join("weclone", "utils", "length_cdf.py")

            command_parts = [
                python_executable,
                script_path,
                f'--model_name_or_path="{self.c["model_name_or_path"]}"',
                f'--dataset="{self.c["dataset"]}"',
                f'--dataset_dir="{self.c["dataset_dir"]}"',
                f'--template="{self.c["template"]}"',
                f"--interval={self.c['cutoff_len']}",
            ]

            child_env = os.environ.copy()
            child_env["CUDA_VISIBLE_DEVICES"] = "0"
            child_env["LLAMAFACTORY_VERBOSITY"] = "ERROR"

            process = subprocess.Popen(
                command_parts,
                env=child_env,
                stdout=None,  # 使用 None 表示使用父进程的标准输出（即终端）
                stderr=None,  # 使用 None 表示使用父进程的标准错误（即终端）
                text=True,
                bufsize=1,  # 行缓冲
            )
            return_code = process.wait()
            if return_code != 0:
                logger.error(f"命令 '{' '.join(command_parts)}' 执行失败，返回码 {return_code}")
        except FileNotFoundError:
            # command_parts[0] 是 python_executable, command_parts[1] 是 script_path
            logger.error(f"命令执行失败: 找不到可执行文件 '{command_parts[0]}' 或脚本 '{command_parts[1]}'")
        except KeyError as e:
            logger.error(f"执行 length_cdf.py 脚本失败：配置项缺失 {str(e)}")
        except Exception as e:
            logger.error(f"执行 length_cdf.py 脚本时发生未知错误: {str(e)}")

    def get_csv_files(self):
        """遍历文件夹获取所有CSV文件路径"""

        csv_files = []
        for chat_obj_folder in os.listdir(self.csv_folder):
            chat_obj_folder_path = os.path.join(self.csv_folder, chat_obj_folder)
            for csvfile in os.listdir(chat_obj_folder_path):
                if not csvfile.endswith(".csv"):
                    continue
                csvfile_path = os.path.join(chat_obj_folder_path, csvfile)
                csv_files.append(csvfile_path)
        return csv_files

    def match_qa(self, messages: List[ChatMessage]) -> List[Union[QaPair, CutMessage]]:
        """
        匹配问答对

        Args:
            messages: 消息列表

        Returns:
            List[Union[QaPair, CutMessage]]: 包含指令和输出的问答对列表
        """
        # 状态定义
        WAITING_INSTRUCTION = "waiting_instruction"  # 等待指令
        WAITING_RESPONSE = "waiting_response"  # 等待回复

        current_state = WAITING_INSTRUCTION
        qa_res: List[Union[QaPair, CutMessage]] = []
        last_message = None
        current_instruction = None
        qa_id_counter = 0

        for msg in messages:
            if isinstance(msg, CutMessage):
                current_state = WAITING_INSTRUCTION
                current_instruction = None
                last_message = None
                if self.c["prompt_with_history"]:
                    qa_res.append(msg)
                continue

            if current_state == WAITING_INSTRUCTION:
                if msg.is_sender == 0:  # 收到对方消息
                    current_instruction = msg.msg
                    last_message = msg
                    current_state = WAITING_RESPONSE

            elif current_state == WAITING_RESPONSE:
                if msg.is_sender == 0:  # 收到对方消息
                    current_instruction = msg.msg
                    last_message = msg
                    # 状态保持不变
                else:  # 自己的回复 使用策略判断是否属于同一对话
                    if last_message and self.qa_match_strategy.is_same_conversation([last_message], msg):
                        assert current_instruction is not None, (
                            "current_instruction should not be None when creating a QA pair"
                        )
                        qa_pair = QaPair(
                            id=qa_id_counter,
                            system=self.system_prompt,
                            instruction=current_instruction,
                            output=msg.msg,
                            history=[],  # No history in this context yet
                            time=msg.CreateTime,  # Use the response message time
                            score=0,  # Default score
                        )
                        qa_res.append(qa_pair)
                        qa_id_counter += 1  # 增加计数器
                    else:
                        if self.c["prompt_with_history"]:
                            qa_res.append(
                                CutMessage(
                                    is_sender=msg.is_sender,
                                    cut_type=msg.type_name,
                                    CreateTime=msg.CreateTime,
                                )
                            )
                    # 无论是否匹配，都重置状态
                    current_state = WAITING_INSTRUCTION
                    current_instruction = None
                    last_message = None

        return qa_res

    # TODO: need review
    def add_history_to_qa(self, qa_res: List[Union[QaPair, CutMessage]]) -> List[QaPair]:
        """
        Adds conversation history to QaPair objects.

        Args:
            qa_res: A list containing QaPair and CutMessage objects.

        Returns:
            A list of QaPair objects with history populated.
        """
        qa_res_with_history: List[QaPair] = []
        current_history: List[List[str]] = []
        last_timestamp: Timestamp = None  # type: ignore

        for item in qa_res:
            if isinstance(item, CutMessage):
                if current_history:
                    instruction = current_history[-1][0]
                    output = current_history[-1][1]
                    history = current_history[:-1]
                    qa_pair_with_history = QaPair(
                        id=-1,
                        system=self.system_prompt,
                        instruction=instruction,
                        output=output,
                        history=history,
                        time=last_timestamp,
                        score=0,
                    )
                    qa_res_with_history.append(qa_pair_with_history)
                current_history = []
                last_timestamp = None  # type: ignore
            elif isinstance(item, QaPair):
                current_history.append([item.instruction, item.output])
                last_timestamp = item.time

        if current_history:
            instruction = current_history[-1][0]
            output = current_history[-1][1]
            history = current_history[:-1]
            # Ensure last_timestamp is not None before assignment
            final_timestamp_end = last_timestamp
            assert final_timestamp_end is not None, "Timestamp cannot be None for the final QaPair"
            qa_pair_with_history = QaPair(
                id=-1,
                system=self.system_prompt,
                instruction=instruction,
                output=output,
                history=history,
                time=final_timestamp_end,
                score=0,
            )
            qa_res_with_history.append(qa_pair_with_history)

        return qa_res_with_history

    def group_consecutive_messages(self, messages: List[ChatMessage]) -> List[ChatMessage]:
        """
        将同一个人连续发送的多条消息组合成一条消息，遇到cut_type添加cut

        Args:
            messages: 消息列表

        Returns:
            List[ChatMessage]: 组合后的消息列表
        """
        if not messages:
            return []

        def _combine_text(messages: List[ChatMessage]) -> ChatMessage:
            """
            合并多条消息为一条

            Args:
                messages: 要合并的消息列表

            Returns:
                ChatMessage: 合并后的消息
            """
            base_msg = messages[0]
            combined_content = messages[0].msg

            for i in messages[1:]:
                content = i.msg
                if not content:
                    continue

                if combined_content and combined_content[-1] not in ["。", "！", "？", "…", "，", "."]:
                    combined_content += "，"

                combined_content += content
            if len(combined_content) > self.c["combine_msg_max_length"]:
                logger.warning(f"组合后消息长度超过{self.c['combine_msg_max_length']}将截断：\n {combined_content}")
                combined_content = combined_content[: self.c["combine_msg_max_length"]]

            combined_message = ChatMessage(
                id=base_msg.id,
                MsgSvrID=base_msg.MsgSvrID,
                type_name=base_msg.type_name,
                is_sender=base_msg.is_sender,
                talker=base_msg.talker,
                room_name=base_msg.room_name,
                msg=combined_content,
                src=base_msg.src,
                CreateTime=messages[-1].CreateTime,  # 使用最后一条消息的时间
            )

            return combined_message

        def _create_cut_message(message: ChatMessage) -> CutMessage:
            return CutMessage(
                is_sender=message.is_sender,
                cut_type=message.type_name,
                CreateTime=message.CreateTime,
            )

        def _combine_current_group(group):
            """
            处理当前消息组并添加到grouped_messages

            Args:
                group: 当前消息组
            """
            if len(group) > 1:
                combined_msg = _combine_text(group)
                grouped_messages.append(combined_msg)
            else:
                grouped_messages.append(group[0])

        grouped_messages = []
        current_group = []

        for _, current_msg in enumerate(messages):
            if current_msg.type_name in self.cut_type_list:
                if current_group:
                    # 当前组有消息，合并当前组，并添加一条cut
                    _combine_current_group(current_group)
                    current_group = []

                    cut_msg = _create_cut_message(current_msg)
                    grouped_messages.append(cut_msg)
                else:
                    # 当前组没消息，检查上一个组
                    if grouped_messages:
                        if not isinstance(grouped_messages[-1], CutMessage):
                            cut_msg = _create_cut_message(current_msg)
                            grouped_messages.append(cut_msg)
                    # 如果上一个组没消息或最后一条是CutMessage，直接continue
                continue

            if not current_group:
                current_group = [current_msg]
                continue

            last_msg = current_group[-1]

            # 判断是否是同一个人的连续消息
            if (
                current_msg.is_sender == last_msg.is_sender
                and current_msg.talker == last_msg.talker
                and self.single_combine_strategy.is_same_conversation([last_msg], current_msg)
            ):
                current_group.append(current_msg)
            else:
                # 不是同一个人的消息，处理当前组并开始新组
                _combine_current_group(current_group)
                # 开始新组
                current_group = [current_msg]

        # 处理最后一组消息
        if current_group:
            _combine_current_group(current_group)

        return grouped_messages

    def process_by_msgtype(self, chat_message: ChatMessage):
        if chat_message.type_name == "文本":
            self.process_text(chat_message)
        # elif chat_message.type_name == "图片":
        #     self.process_image(chat_message)

    def load_csv(self, file_path) -> List[ChatMessage]:
        """
        做整体第一次预处理，过滤不符合条件的行
        """
        df = pd.read_csv(file_path, encoding="utf-8", dtype={"msg": str})

        df = df[~df["type_name"].isin(values=skip_type_list)]

        # 如果type_name为文本 并且msg 包含 手机号、身份证号、邮箱、网址则删除这行
        for i in df.index:
            if df.loc[i, "type_name"] == "文本":
                msg_str = str(df.loc[i, "msg"])
                if (
                    re.search(r"1\d{10}", msg_str)
                    or re.search(r"\d{18}", msg_str)
                    or re.search(r"\w+@\w+", msg_str)
                    or "http" in msg_str
                    or r"\\xa0" in msg_str
                    or r"\\u" in msg_str
                ):
                    df = df.drop(index=i)
                    continue
                for blocked_word in self.blocked_words:
                    if blocked_word in msg_str:
                        df = df.drop(index=i)
                        break
            else:
                df.loc[i, "msg"] = ""

        df = df.dropna(how="all")
        # 时间格式 2021-07-07 10:27:23
        # 遍历行 相同is_sender的行合并msg（）遇到不同is_sender就重新开始
        df["CreateTime"] = pd.to_datetime(df["CreateTime"])

        return [ChatMessage(*row) for row in df.values]

    def process_text(self, chat_message: ChatMessage):
        pass

    def save_result(self, qa_res: List[QaPair]):
        """
        Saves the list of QaPair objects to a JSON file after converting them to dictionaries.

        Args:
            qa_res: A list of QaPair objects.
        """
        processed_qa_res = []
        for idx, item in enumerate(qa_res):
            item_dict = {
                "id": idx,
                "system": item.system,
                "instruction": item.instruction,
                "output": item.output,
                "history": item.history,
                "time": item.time.isoformat() if item.time else None,
                "score": item.score,
            }
            processed_qa_res.append(item_dict)

        output_path = "./dataset/res_csv/sft/sft-my.json"
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(processed_qa_res, f, ensure_ascii=False, indent=4)
        logger.success(f"聊天记录处理成功，共{len(qa_res)}条，保存到 {output_path}")


if __name__ == "__main__":
    processor = DataProcessor()
    processor.main()
